#include "kalman.h"
#include <math.h>


// printing on each build system
#ifdef HAVE_PRINT
#include <stdio.h>
#define ERROR printf("FAIL\n")
#else
// have to pass around the USART1 pointer to get this going??
#define ERROR
#endif

void matvecmul_add_ge_3x1(kal_t *A, kal_t *x, kal_t *y, kal_t *z) {
	/*
	 x 1x1 vector
	 A 3x1 general
	 y 3x1 vector
	 z 3x1 vector 
	 z = z + y + Ax
	 */
	
	z[0] = z[0] + y[0] + A[0]*x[0];
	z[1] = z[1] + y[1] + A[1]*x[0];
	z[2] = z[2] + y[2] + A[2]*x[0];
}

void matvecmul_tr_3x3_ip(kal_t *A, kal_t *x) {
	/*
	 x 3x1 vector
	 A 3x3 lower unit triangular 
	 x = Ax
	 */
	
	kal_t x1, x2;
	x1 = x[1];
	x2 = x[2];
	
	x[1] = A[1]*x[0] + x1;
	x[2] = A[2]*x[0] + A[5]*x1  + x2;
}

void matvecmul_ge_3x3_ip(kal_t *A, kal_t *x) {
	/*
	 x 3x1 vector
	 A 3x3 general 
	 x = Ax
	 */
	
	kal_t x0, x1, x2;
	x0 = x[0];
	x1 = x[1];
	x2 = x[2];
	
	x[0] = A[0]*x0 + A[3]*x1 + A[6]*x2;
	x[1] = A[1]*x0 + A[4]*x1 + A[7]*x2;
	x[2] = A[2]*x0 + A[5]*x1 + A[8]*x2;
}

void vecsub_3x1_ip(kal_t *a, kal_t *b) {
	/*
	 a = a - b
	 */
	a[0] = a[0] - b[0];
	a[1] = a[1] - b[1];
	a[2] = a[2] - b[2];
}

void matadd_sub_ge_3x3_ip(kal_t *a, kal_t *b, kal_t *c) {
	/*
	 a = a + b - c
	 */
	a[0] = a[0] + b[0] - c[0];
	a[1] = a[1] + b[1] - c[1];
	a[2] = a[2] + b[2] - c[2];
	
	a[3] = a[3] + b[3] - c[3];
	a[4] = a[4] + b[4] - c[4];
	a[5] = a[5] + b[5] - c[5];
	
	a[6]  = a[6]  + b[6]  - c[6];
	a[7]  = a[7]  + b[7]  - c[7];
	a[8]  = a[8]  + b[8]  - c[8];
}

void matadd_ge_3x3(kal_t *a, kal_t *b, kal_t *c) {
	/*
	 c = a + b
	 */
	c[0] = a[0] + b[0];
	c[1] = a[1] + b[1];
	c[2] = a[2] + b[2];
	
	c[3] = a[3] + b[3];
	c[4] = a[4] + b[4];
	c[5] = a[5] + b[5];
	
	c[6]  = a[6]  + b[6];
	c[7]  = a[7]  + b[7];
	c[8]  = a[8]  + b[8];
}

void matmul_ge_NN_3x3_ip(kal_t *a, kal_t *b) {
	/*
	 in place matrix multiply
	 b = a * b
	 a 3x3 general
	 b 3x3 general
	 */
	
	kal_t b0, b1, b2;
	
	b0 = b[0];
	b1 = b[1];
	b2 = b[2];
	
	b[0] = a[0]*b0 + a[3]*b1 + a[6]*b2;
	b[1] = a[1]*b0 + a[4]*b1 + a[7]*b2;
	b[2] = a[2]*b0 + a[5]*b1 + a[8]*b2;
	
	b0 = b[3];
	b1 = b[4];
	b2 = b[5];
	
	b[3] = a[0]*b0 + a[3]*b1 + a[6]*b2;
	b[4] = a[1]*b0 + a[4]*b1 + a[7]*b2;
	b[5] = a[2]*b0 + a[5]*b1 + a[8]*b2;

	b0 = b[6];
	b1 = b[7];
	b2 = b[8];
	
	b[6] = a[0]*b0 + a[3]*b1 + a[6]*b2;
	b[7] = a[1]*b0 + a[4]*b1 + a[7]*b2;
	b[8] = a[2]*b0 + a[5]*b1 + a[8]*b2;
}

void matmul_tr_NN_3x3(kal_t *a, kal_t *b, kal_t *c) {
	/*
	 c = a * b
	 c 3x3 general
	 a 3x3 unit lower triangular
	 b 3x3 general
	 
	 used for A(PC') and A(PA')
	 */
	
	c[0] = b[0];
	c[1] = a[1]*b[0] + b[1];
	c[2] = a[2]*b[0] + a[5]*b[1] + b[2];
	
	c[3]  = b[3];
	c[4]  = a[1]*b[3] + b[4];
	c[5]  = a[2]*b[3] + a[5]*b[4] + b[5];
	
	c[6] = b[6];
	c[7] = a[1]*b[6] + b[7];
	c[8] = a[2]*b[6] + a[5]*b[7] + b[7];
}

void matmul_tr_NT_3x3(kal_t *a, kal_t *b, kal_t *c) {
	/*
	 c = a * b'
	 c 3x3 general
	 a 3x3 general
	 b 3x3 unit lower triangular - transpose
	 
	 used for PA'
	 */
	
	c[0] = a[0];
	c[1] = a[1];
	c[2] = a[2];
	
	c[3]  = a[0]*b[1] + a[3];
	c[4]  = a[1]*b[1] + a[4];
	c[5]  = a[2]*b[1] + a[5];
	
	c[6] = a[0]*b[2] + a[3]*b[5]  + a[6];
	c[7] = a[1]*b[2] + a[4]*b[5]  + a[7];
	c[8] = a[2]*b[2] + a[5]*b[5]  + a[8];
}

void solve(kal_t *A, kal_t *B) {
	/*
	 A 3x3 general matrix, input (APC')/output(K)
	 B 3x3 general matrix, input (CPC'+Sz) - destroyed
	 
	 solve K = APC'(CPC'+Sz)^(-1) for K.
	 rewrite as
	 A = APC' and B = (CPC'+Sz)
	 K = AB^(-1)
	 then as
	 KB = A
	 
	 1) perform gaussian elimination with column pivoting on B 
	 leaving lower triangular
	 
	 2) back solve for K, inplace (A)

	 -----------------------
	 to check, at end K == AB^(-1)
	*/
	
	int i, j, k;

	// start factorization
	for (i = 0; i < 3; ++i) {
		kal_t *currLoc = B + i*3 + i;
		
		// determine if pivot is necessary
		int piv = i;
		kal_t pivMax = currLoc[0];
		currLoc += 3;
		for (j = i+1; j < 3; ++j) {
			if (fabs(currLoc[0]) > pivMax) {
				piv = j;
				pivMax = fabs(currLoc[0]);
			}
			currLoc += 3;
		}
		
		// pivot when appropriate
		if (i != piv) {
			kal_t *Ac0 = A + i*3;
			kal_t *Ac1 = A + piv*3;
			
			kal_t *Bc0 = B + i*3;
			kal_t *Bc1 = B + piv*3;
			
			for (j = 0; j < 3; ++j) {
				kal_t tmp = Ac0[j];
				Ac0[j] = Ac1[j];
				Ac1[j] = tmp;
				
				tmp = Bc0[j];
				Bc0[j] = Bc1[j];
				Bc1[j] = tmp;
			}
		}
		
		// calculate scales and store inplace
		kal_t *curr = B + i;
		
		// got a problem here, not sure how we want to deal with this
		if (curr[i*3] == 0.0) {
			ERROR
		}
		
		for (j = i+1; j < 3; ++j) {
			curr[j*3] = curr[j*3]/curr[i*3];
		}
		
		// update B
		kal_t *fact = B + i;
		for (j = i+1; j < 3; ++j) {
			curr = B + j;
			for (k = i+1; k < 3; ++k) {
				curr[k*3] -= curr[i*3]*fact[k*3];
			}
		}
		
		// update A
		fact = B + i;
		for (j = 0; j < 3; ++j) {
			curr = A + j;
			for (k = i+1; k < 3; ++k) {
				curr[k*3] -= curr[i*3]*fact[k*3];
			}
		}
		
//#define DEBUG		
#ifdef DEBUG
		// DEBUG only
		// zero lower triangular
		curr = B + i;
		for (j = i+1; j < 3; ++j) {
			curr[j*3] = 0;
		}
#endif
	}
	// end factorization
	
	// perform back solve for each row of A
	for (k = 0; k < 3; ++k) {
		kal_t *Ar0 = A + k;
		
		// for each column of B
		for (i = 2; i >= 0; --i) {
			kal_t *Bc0 = B + i*3;
			kal_t acc0 = 0.0;
			
			// for the elemnts below the diagonal
			for (j = 2; j > i; --j) {
				acc0 += Ar0[j*3] * Bc0[j];
			}
			if (Bc0[i] == 0.0) {
				ERROR;
			}
			Ar0[i*3] = (Ar0[i*3] - acc0) / Bc0[i];
		}
	}
}


void kalman(kal_t *P, kal_t *K, kal_t *y, kal_t *x, kal_t *u, int update) {
	
	// time dependent constants
	static kal_t A[9]      // 3x3 relationship between state variables
			= {1.0, time_v, 0.5*time_v*time_v, 
			   0.0, 1.0,    time_v,     
			   0.0, 0.0,    1.0 };
	
	static kal_t B[3]      // 3x1 relationsiop between commanded input and state variables
			= {1.0, time_v, 0.5*time_v*time_v };
	
	
	// constants  
	//kal_t C[9];			// relationship between sensor measurements and expected state variables
							// C is currently identity and not needed
	static kal_t Sz[9]      // 3x3 noise caused by instrument error 
							//    (large values means do not give much credibility to measurements)
			= {	0.09649*1.5,    0.09572*1.5,    0.01419*1.5,
				0.01576*1.5,    0.04854*1.5,    0.04218*1.5,
				0.09706*1.5,    0.08003*1.5,    0.09157*1.5 };
	static kal_t Szi[9]    // 3x3 inverse of Sz
			= { 0.03473*1.5,   -.24778*1.5,    .10874*1.5,
			    0.08607*1.5,    .24223*1.5,   -.12490*1.5,
			   -.11203*1.5,    0.05093*1.5,    .10310*1.5 };
	static kal_t Sw[9]		// 3x3 process noise
			= {	0.0383*.25,    0.0631*.25,    0.0203*.25,
				0.0631*.25,    0.1041*.25,    0.0334*.25,
				0.0203*.25,    0.0334*.25,    0.0107*.25 };
	
	// temporary data structures
	kal_t tmp[9];
	kal_t *PCt;
	kal_t PAt[9];

	
	// compute PAt and PCt
	matmul_tr_NT_3x3(P,A,PAt);

	// C is currently identity matrix so PCt = P
	// done with this iterations use of P so just
	// pointing PCt to P
	PCt = P;

	// kalman gain
	// K = APC'(CPC'+Sz)^-1 rewritten as K(CPC'+Sz) = APC'
	matmul_tr_NN_3x3(A,PCt,K);
	// C is currently identity 
	matadd_ge_3x3(Sz,PCt,tmp);
	solve(K,tmp);
	
	// next state
	//x = (Ax + Bu) + K(y-Cx)
	// C is currently identity matrix
	if (update) {
		vecsub_3x1_ip(y,x);
		matvecmul_ge_3x3_ip(K,y);
		matvecmul_tr_3x3_ip(A,x);
		matvecmul_add_ge_3x1(B,u,y,x);
	}
	else {
		//vecsub_3x1_ip(y,x);
		//matvecmul_ge_3x3_ip(K,y);
		kal_t l[3];
		l[0] = 0.0;
		l[1] = 0.0;
		l[2] = 0.0;
		matvecmul_tr_3x3_ip(A,x);
		matvecmul_add_ge_3x1(B,u,l,x);
	}

	// calculate error covariance
	// P = APA' + Sw - APC'(Szi)CPA'
	// NOTE: PCt points to P
	matmul_tr_NN_3x3(A,PCt,tmp);
	// C is currently identity
	matmul_tr_NN_3x3(A,PAt,P);
	matmul_ge_NN_3x3_ip(Szi,PAt);
	matmul_ge_NN_3x3_ip(tmp,PAt);
	matadd_sub_ge_3x3_ip(P,Sw,PAt);
}


#ifdef HAVE_PRINT
void printMatrix(kal_t *A) {
	int i, j;
	for (i = 0; i < 3; ++i) {
		for (j = 0; j < 3; ++j) {
			printf("%f  ",A[j*3+i]);
		}
		printf("\n");
	}
	printf("\n");
}

void printVector(kal_t *x) {
	int i;
	for (i = 0; i < 3; ++i) 
		printf("%f  ",x[i]);
	printf("\n");
}

#endif

/*
 K = AP inv(P + S)
 rewrites to
 K(P+S) = AP
 
 do gaussian elimination using columns of (P+S)
 
 for one row of unknow K, (P+S) will be lower triangular
 and will have update one row of AP
			 s 0 0 0 0 0
			 s s 0 0 0 0
			 s s s 0 0 0
			 s s s s 0 0
			 s s s s s 0
			 s s s s s s
 
 k k k k k k a a a a a a
 
 allowing for backsolve of the current K row.
 do this backsolve for each k row.

*/












